package com.bw.test2;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.dataproc.vector.VectorAssemblerBatchOp;
import com.alibaba.alink.operator.batch.evaluation.EvalRegressionBatchOp;
import com.alibaba.alink.operator.batch.feature.OneHotPredictBatchOp;
import com.alibaba.alink.operator.batch.feature.OneHotTrainBatchOp;
import com.alibaba.alink.operator.batch.regression.LinearRegPredictBatchOp;
import com.alibaba.alink.operator.batch.regression.LinearRegTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.RegressionMetrics;
import org.apache.flink.types.Row;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class OnHotTest {
    @Test
    public void testLinearRegTrainBatchOp() throws Exception {

        BatchOperator.setParallelism(1);
        // 1.准备数据源
        List<Row> df = Arrays.asList(
                Row.of(2, "a", 1),
                Row.of(3, "b", 1),
                Row.of(4, "c", 2),
                Row.of(2, "d", 1),
                Row.of(2, "a", 1),
                Row.of(4, "b", 2),
                Row.of(1, "b", 1),
                Row.of(5, "c", 3)
        );


        // 2、选择特征和label
        BatchOperator<?> batchData = new MemSourceBatchOp(df, "f0 int, f1 string, label int");


        // 独热编码列
        BatchOperator <?> one_hot = new OneHotTrainBatchOp().setSelectedCols("f1");


        BatchOperator <?> one_hot_model = batchData.link(one_hot);

        one_hot_model.lazyPrint(10);

        // 输出的编码列
        BatchOperator <?> one_hot_predictor = new OneHotPredictBatchOp().setOutputCols("f11");

        // 编码后
        BatchOperator<?> one_hot_batch_data = one_hot_predictor.linkFrom(one_hot_model, batchData);

        // 向量聚合
        VectorAssemblerBatchOp vectorAssemblerBatchOp = new VectorAssemblerBatchOp().setSelectedCols("f0", "f11")
                .setOutputCol("vec").linkFrom(one_hot_batch_data);


        // 训练组件
        BatchOperator<?> lr = new LinearRegTrainBatchOp()
//                .setFeatureCols("f0", "f11")
                .setVectorCol("vec")
                .setLabelCol("label")
                .lazyPrintModelInfo();

        // 3. 训练模型
        BatchOperator model = vectorAssemblerBatchOp.link(lr);

        // 4、预测
        // 预测组件
        BatchOperator<?> predictor = new LinearRegPredictBatchOp()
                .setPredictionCol("pred");

        BatchOperator<?> result = predictor.linkFrom(model, vectorAssemblerBatchOp);



        // 评估指标
        RegressionMetrics metrics = new EvalRegressionBatchOp().setPredictionCol("pred").setLabelCol("label").linkFrom(
                result).collectMetrics();
        System.out.println("Total Samples Number:" + metrics.getCount());
        System.out.println("SSE:" + metrics.getSse());
        System.out.println("SAE:" + metrics.getSae());
        System.out.println("RMSE:" + metrics.getRmse());
        System.out.println("R2:" + metrics.getR2());
        System.out.println("MSE:" + metrics.getMse());
    }
}